package com.webank.wedatasphere.linkis.hadoop.common.utils

import java.io.File
import java.nio.file.Paths

import com.webank.wedatasphere.linkis.common.conf.Configuration.hadoopConfDir
import com.webank.wedatasphere.linkis.hadoop.common.conf.HadoopConf._
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.Path
import org.apache.hadoop.security.UserGroupInformation
import org.slf4j.LoggerFactory

private[utils] object UGIUtils {

  val log=LoggerFactory.getLogger("UGIUtils")

  private var ugi: UserGroupInformation = _

  def getConfiguration: Configuration = getConfiguration(hadoopConfDir)

  private def getConfiguration(hadoopConfDir: String): Configuration = {
    val configuration = new Configuration()
    fillConfiguration(configuration,hadoopConfDir)
    configuration
  }
  def fillConfiguration(conf: Configuration): Unit = {
    fillConfiguration(conf,hadoopConfDir)
  }
  private def fillConfiguration(conf: Configuration,hadoopConfDir: String): Unit = {
    val confPath = new File(hadoopConfDir)
    if(!confPath.exists() || confPath.isFile) {
      throw new RuntimeException(s"Create hadoop configuration failed, path $hadoopConfDir not exists.")
    }
    addFileResource(hadoopConfDir,"core-site.xml",conf)
    addFileResource(hadoopConfDir,"hdfs-site.xml",conf)
    addFileResource(hadoopConfDir,"yarn-site.xml",conf)
    addFileResource(hadoopConfDir,"hbase-site.xml",conf)
  }
  private def addFileResource(hadoopConfDir:String,fileName:String,conf:Configuration):Unit = {
    val file = Paths.get(hadoopConfDir, fileName).toAbsolutePath.toFile
    if (file.exists()) {
      conf.addResource(new Path(file.getAbsolutePath))
    }
  }
  def getUGI: UserGroupInformation = {
    ugi.checkTGTAndReloginFromKeytab()
    ugi
  }
  private def getKerberosUser(userName: String): String = {
    var user = userName
    if(KEYTAB_HOST_ENABLED.getValue){
      user = user+ "/" + KEYTAB_HOST.getValue
    }
    user
  }

  val hadoopUser: String = HADOOP_ROOT_USER.getValue
  if(KERBEROS_ENABLE.getValue) {
    val path = new File(KEYTAB_FILE.getValue).getPath
    val user = getKerberosUser(hadoopUser)
    UserGroupInformation.setConfiguration(getConfiguration)
    UserGroupInformation.loginUserFromKeytab(user,path)
    ugi = UserGroupInformation.loginUserFromKeytabAndReturnUGI(user, path)
    log.info("安全认证登录成功")
  } else {
    ugi = UserGroupInformation.createRemoteUser(hadoopUser)
  }
}